#include "stdafx.h"
#include <Psapi.h>


#define IMAGE_VIRTUAL_TO_ABOSLUTE(type, hModule, virtualAdress) ((type)(DWORD(hModule) + DWORD(virtualAdress)))
#define IMPORT_DIRECTORY(hModule) IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_IMPORT_DESCRIPTOR, hModule, GetDataDirectory(hModule)[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress)
#define EXPORT_DIRECTORY(hModule) IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_EXPORT_DIRECTORY, hModule, GetDataDirectory(hModule)[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress)

BOOL IsStringsEqualIgnoreCase(LPCSTR strFirst, LPCSTR strSecond)
{
	char szFirst[1024], szSecond[1024];
	strcpy(szFirst, strFirst); strcpy(szSecond, strSecond);
	strupr(szFirst); strupr(szSecond);
	return strcmp(szFirst, szSecond) == 0;
}

PIMAGE_DATA_DIRECTORY GetDataDirectory(HMODULE hModule)
{
	PIMAGE_DOS_HEADER pimage_dos_header = PIMAGE_DOS_HEADER(hModule);
	if (pimage_dos_header->e_magic != IMAGE_DOS_SIGNATURE)
		throw;

	PIMAGE_NT_HEADERS pimage_nt_headers = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_NT_HEADERS, hModule, pimage_dos_header->e_lfanew);
	if (pimage_nt_headers->Signature != IMAGE_NT_SIGNATURE)
		throw;

	return pimage_nt_headers->OptionalHeader.DataDirectory;
}

void WriteToProtectedMemory(LPDWORD lpTarget, DWORD value)
{
	DWORD dwOldVAlue;
	VirtualProtect(lpTarget, sizeof(DWORD), PAGE_READWRITE, &dwOldVAlue);
	*lpTarget = value;
	VirtualProtect(lpTarget, sizeof(DWORD), dwOldVAlue, &dwOldVAlue);
}

/*void HookExportTable(HMODULE hModule, int iNumberOfProcs, LPCSTR astrProcNames[], LPVOID alpHookProc[], LPVOID alpOriginalProc[])
{
	PIMAGE_EXPORT_DIRECTORY pExportDirectory = EXPORT_DIRECTORY(hModule);

	LPDWORD lpNames = IMAGE_VIRTUAL_TO_ABOSLUTE(LPDWORD, hModule, pExportDirectory->AddressOfNames);
	for (int i = 0; i < pExportDirectory->NumberOfNames; i++)
		for (int j = 0; j < iNumberOfProcs; j++)
			if (IsStringsEqualIgnoreCase(IMAGE_VIRTUAL_TO_ABOSLUTE(LPCSTR, hModule, lpNames[i]), astrProcNames[j]))
			{
				LPDWORD lpEntry = (IMAGE_VIRTUAL_TO_ABOSLUTE(LPDWORD, hModule, pExportDirectory->AddressOfFunctions)) + i;
				alpOriginalProc[j] = IMAGE_VIRTUAL_TO_ABOSLUTE(LPVOID, hModule, *lpEntry);
				WriteToProtectedMemory(lpEntry, (DWORD)alpHookProc[j]);
			}
}*/

void HookImportTable(HMODULE hModule, LPCSTR strImportName, int iNumberOfProcs, LPCSTR astrProcNames[], LPVOID alpHookProc[], LPVOID alpOriginalProc[])
{
	IMAGE_DATA_DIRECTORY importDirectory = GetDataDirectory(hModule)[IMAGE_DIRECTORY_ENTRY_IMPORT];

	if (importDirectory.Size > 0)
	{
		PIMAGE_IMPORT_DESCRIPTOR lpImports = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_IMPORT_DESCRIPTOR, hModule, importDirectory.VirtualAddress);

		while(lpImports->Name != 0)
		{
			LPCSTR strModuleName = IMAGE_VIRTUAL_TO_ABOSLUTE(LPCSTR, hModule, lpImports->Name);

			if (IsStringsEqualIgnoreCase(strImportName, strModuleName))
			{
				PIMAGE_THUNK_DATA pOriginalThunk = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_THUNK_DATA, hModule, lpImports->OriginalFirstThunk);
				PIMAGE_THUNK_DATA pThunk = IMAGE_VIRTUAL_TO_ABOSLUTE(PIMAGE_THUNK_DATA, hModule, lpImports->FirstThunk);

				while(pOriginalThunk->u1.AddressOfData != 0)
				{
					if (pOriginalThunk->u1.Ordinal & 0x80000000)
					{
						DWORD dwOrdinal = pOriginalThunk->u1.Ordinal & 0x7FFFFFFF;
					}
					else
					{
						LPWORD lpHint = IMAGE_VIRTUAL_TO_ABOSLUTE(LPWORD, hModule, pOriginalThunk->u1.AddressOfData);
						LPCSTR lpFunctionName = IMAGE_VIRTUAL_TO_ABOSLUTE(LPCSTR, hModule, pOriginalThunk->u1.Function + 2);

						for (int i = 0; i < iNumberOfProcs; i++)
							if (IsStringsEqualIgnoreCase(astrProcNames[i], lpFunctionName))
							{
								alpOriginalProc[i] = (LPVOID)pThunk->u1.AddressOfData;
								WriteToProtectedMemory((LPDWORD)pThunk, (DWORD)alpHookProc[i]);
							}
					}

					pOriginalThunk++;
					pThunk++;
				}
			}

			lpImports++;
		}
	}
}

void Hook(LPCSTR strModule, int iNumberOfProcs, LPCSTR astrProcNames[], LPVOID alpHookProc[], LPVOID alpOriginalProc[])
{
	HMODULE hModuleToBeHooked = LoadLibraryA(strModule);
	//HookExportTable(hModuleToBeHooked, iNumberOfProcs, astrProcNames, alpHookProc, alpOriginalProc);

	for (int i = 0; i < iNumberOfProcs; i++)
		alpOriginalProc[i] = GetProcAddress(hModuleToBeHooked, astrProcNames[i]);

	HMODULE modules[1024];
	DWORD cbNeeded;
	EnumProcessModules(GetCurrentProcess(), modules, sizeof(modules), &cbNeeded);
	DWORD dwModulesCount = cbNeeded / 4;

	LPVOID* dummy = new LPVOID[iNumberOfProcs];
	for (int i = 0; i < dwModulesCount; i++)
		if (modules[i] != hModuleToBeHooked)
			HookImportTable(modules[i], strModule, iNumberOfProcs, astrProcNames, alpHookProc, dummy);
	delete[] dummy;

	FreeLibrary(hModuleToBeHooked);
}